import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from dataset import MyDataset
from utils import read_split_data
from GANs.model import Generator,Discriminator

if __name__=='__main__':
    args = parsers()
    train_images_path, train_images_label, _, _ = read_split_data(args.data_path)
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                                   batch_size = batch_size,
                                                   shuffle = True,
                                                   pin_memory = False,
                                                   num_workers = 0,
                                                   collate_fn = train_dataset.collate_fn)
    generator = Generator()
    discriminator = Discriminator()
    optimizer_G = optim.Adam(generator.parameters(), lr=0.0002)
    optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002)
    criterion = nn.BCELoss()

    epochs = 5
    for epoch in range(epochs):
        for i, (images, _) in enumerate(train_loader):
            # 真实图像标签是1，生成图像标签是0
            real_labels = torch.ones(images.size(0), 1)
            fake_labels = torch.zeros(images.size(0), 1)
            # 训练判别器
            outputs = discriminator(images)
            d_loss_real = criterion(outputs, real_labels)
            real_score = outputs
            z = torch.randn(images.size(0), 100)
            fake_images = generator(z)
            outputs = discriminator(fake_images.detach())
            d_loss_fake = criterion(outputs, fake_labels)
            fake_score = outputs
            d_loss = d_loss_real + d_loss_fake
            optimizer_D.zero_grad()
            d_loss.backward()
            optimizer_D.step()
            # 训练生成器
            outputs = discriminator(fake_images)
            g_loss = criterion(outputs, real_labels)
            optimizer_G.zero_grad()
            g_loss.backward()
            optimizer_G.step()
        print(f'Epoch [{epoch+1}/{epochs}], d_loss: {d_loss.item():.4f}, g_loss: {g_loss.item():.4f}')

    z = torch.randn(1, 100)
    generated_images = generator(z)
    plt.imshow(generated_images.detach().numpy().reshape(28, 28), cmap='gray')
    plt.show()

